{-# LANGUAGE ViewPatterns               #-}
{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  HMinCaml.ConstFold
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Constant folding.
--
--------------------------------------------------------------------------------

module HMinCaml.ConstFold
  (

    constFold

  ) where

import qualified HMinCaml.HMonad as H
import HMinCaml.Id
import HMinCaml.KNormal


import Control.Applicative
import Control.Monad
import qualified Data.Map as M
import Data.List ( foldl' )


--------------------------------------------------------------------------------
-- Inline expansion


type Env = M.Map IdS Expr

type Mon a = H.HMonad a

mem :: IdS -> (Expr -> Maybe a) -> Env -> Maybe a
mem x match = match <=< M.lookup x


findi :: IdS -> Env -> Maybe Int
findi x = mem x (\a -> case a of (Int i) -> Just i; _ -> Nothing)

findf :: IdS -> Env -> Maybe Double
findf x = mem x (\a -> case a of (Float d) -> Just d; _ -> Nothing)

findt :: IdS -> Env -> Maybe [IdS]
findt x = mem x (\a -> case a of (Tuple xs) -> Just xs; _ -> Nothing)

{-
findi :: IdS -> Env -> Mon Int
findi x env = case M.lookup x env of
    Just (Int i) -> return i
    _            -> H.fatal "Not found"
-}

constFold :: Expr -> H.HMonad Expr
constFold expr = cf M.empty expr


cf :: Env -> Expr -> Mon Expr
cf env (Var x)            | Just i <- findi x env = return $ Int i
cf env (Neg x)            | Just i <- findi x env = return $ Int $ negate i

cf env (Add x y)          | Just i <- findi x env
                          , Just j <- findi y env 
                          = return $ Int $ i + j

cf env (Sub x y)          | Just i <- findi x env
                          , Just j <- findi y env 
                          = return $ Int $ i - j

cf env (FNeg x)           | Just d <- findf x env = return $ Float $ negate d

cf env (FAdd x y)         | Just s <- findf x env
                          , Just t <- findf y env 
                          = return $ Float $ s + t

cf env (FSub x y)         | Just s <- findf x env
                          , Just t <- findf y env 
                          = return $ Float $ s - t

cf env (FMul x y)         | Just s <- findf x env
                          , Just t <- findf y env 
                          = return $ Float $ s * t

cf env (FDiv x y)         | Just s <- findf x env
                          , Just t <- findf y env 
                          = return $ Float $ s / t


cf env (IfEq x y e1 e2)   | Just i <- findi x env
                          , Just j <- findi y env 
                          = if i == j then cf env e1 else cf env e2


cf env (IfEq x y e1 e2)   | Just s <- findf x env
                          , Just t <- findf y env 
                          = if s == t then cf env e1 else cf env e2

cf env (IfEq x y e1 e2)   = IfEq x y <$> cf env e1 <*> cf env e2

cf env (IfLE x y e1 e2)   | Just i <- findi x env
                          , Just j <- findi y env 
                          = if i <= j then cf env e1 else cf env e2


cf env (IfLE x y e1 e2)   | Just s <- findf x env
                          , Just t <- findf y env 
                          = if s <= t then cf env e1 else cf env e2

cf env (IfLE x y e1 e2)   = IfLE x y <$> cf env e1 <*> cf env e2

cf env (Let (x,t) e1 e2)  = do
    ea <- cf env e1
    eb <- cf (M.insert x ea env) e2
    return $ Let (x,t) ea eb

cf env (LetRec (Fundef xt args e1) e2) = 
   (\ea eb -> LetRec (Fundef xt args ea) eb) <$> cf env e1 <*> cf env e2

cf env (LetTuple xts y e) | Just ys <- findt y env
                          = (\ea -> (foldl' fn ea $ zip xts ys)) <$> cf env e
  where
    fn e' (xt,z) = Let xt (Var z) e'

cf env (LetTuple xts y e) = LetTuple xts y <$> cf env e

cf _   e                  = return e
